import os
import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import argrelextrema
"""
from multiprocessing import cpu_count
from concurrent.futures import ProcessPoolExecutor
"""

save_path = os.getcwd()

m_1 = 20
m_2 = 10
k_1 = 600
k_2 = 1000
last_time = 50

# 设置注释文本的样式与箭头样式
BBOX = dict(boxstyle="round", fc="0.8")
ARROW = dict(
    arrowstyle="->",
    connectionstyle="angle,angleA=90,angleB=0,rad=10")

def Simulation(position, velocity, delta_t, func, moment, sim_method):
    x_1, x_2 = position
    v_1, v_2 = velocity
    a_1 = ((x_2 - x_1) * k_2 - x_1 * k_1) / m_1
    a_2 = ((x_1 - x_2) * k_2 + func(moment)) / m_2
    if sim_method == 'EulerCromer':
        v_1 += a_1 * delta_t
        v_2 += a_2 * delta_t
        x_1 += v_1 * delta_t
        x_2 += v_2 * delta_t
    elif sim_method == 'Euler':
        x_1 += v_1 * delta_t
        x_2 += v_2 * delta_t
        v_1 += a_1 * delta_t
        v_2 += a_2 * delta_t
    elif sim_method == 'EulerRichardson':
        x_1_half = x_1 + v_1 * delta_t / 2
        x_2_half = x_2 + v_2 * delta_t / 2
        v_1_half = v_1 + a_1 * delta_t / 2
        v_2_half = v_2 + a_2 * delta_t / 2
        a_1_half = ((x_2_half - x_1_half) * k_2 - x_1_half * k_1) / m_1
        a_2_half = ((x_1_half - x_2_half) * k_2 + func(moment + delta_t / 2)) / m_2
        x_1 += v_1_half * delta_t
        x_2 += v_2_half * delta_t
        v_1 += a_1_half * delta_t
        v_2 += a_2_half * delta_t
    return (x_1, x_2), (v_1, v_2)

class my_func:
    def __init__(self, omega):
        self.Omega = omega

    def my_sin(self, moment):
        return 5 * np.sin(self.Omega * moment)

    def my_half_sin(self, moment):
        answer = self.my_sin(moment)
        return np.where(answer < 0, 0, answer)

    def get_theta(self, moment):
        return self.Omega * moment - np.floor(self.Omega * moment / (np.pi * 2)) * 2 * np.pi

    def my_triangle(self, moment):
        theta = self.get_theta(moment)
        return np.where(theta < np.pi, 10 * (theta / np.pi - 0.5), 10 * (1.5 - theta / np.pi))

    def my_rectangle(self, moment):
        theta = self.get_theta(moment)
        return np.where(theta < np.pi, 5, -5)

def draw(sim_method, func):
    x_2_max_array = []
    if sim_method == "Eular":
        precision_default = 50000
    else:
        precision_default = 5000
    omega_array = np.linspace(1, 20, 400)
    for omega in omega_array:
        Func = my_func(omega)
        moments = np.linspace(0, last_time, precision_default)
        delta_t = last_time / precision_default
        position = [0, 0]
        x_2_max = position[1]
        velocity = [0, 0]

        for moment in moments[:-1]:
            position, velocity = Simulation(position, velocity, delta_t, getattr(Func, "my_" + func), moment, sim_method)
            x_2_max = max(x_2_max, position[1])
        x_2_max_array.append(x_2_max)

    x_2_max_array = np.array(x_2_max_array)
    x_2_max_greater_arg = np.array(argrelextrema(x_2_max_array, np.greater))[0]
    x_2_max_greater_length = len(x_2_max_greater_arg)
    if x_2_max_greater_length >= 1:
        x_2_max_greater = x_2_max_array[x_2_max_greater_arg]
        x_2_max_greatest_arg = x_2_max_greater_arg[x_2_max_greater.argsort()[-(2 if x_2_max_greater_length > 2 else x_2_max_greater_length):][::-1]]
        plt.scatter(omega_array[x_2_max_greatest_arg], x_2_max_array[x_2_max_greatest_arg])
        for arg in x_2_max_greatest_arg:
            x_2 = x_2_max_array[arg]
            omega = omega_array[arg]
            plt.annotate(
                "$\\omega=%frad\cdot{s}^{-1}$\n${x}_{2,max}=%em$"%(omega, x_2),
                xy=(omega, x_2),
                xytext=(0, 20),
                textcoords='offset points',
                bbox=BBOX,
                arrowprops=ARROW
            )
    plt.plot(omega_array, x_2_max_array)
    plt.xlabel("$\\omega/(rad\cdot{s}^{-1})$")
    plt.ylabel("${x}_{2,max}/m$")
    plt.savefig(os.path.join(save_path, "%s_%s.png"%(func, sim_method)))
    plt.clf()

if __name__ == "__main__":
    '''
    with ProcessPoolExecutor(max_workers=cpu_count()) as executor:
        for sim_method in ["Euler", "EulerCromer", "EulerRichardson"]:
            for func in ["sin", "half_sin", "triangle", "rectangle"]:
                executor.submit(draw, sim_method, func)
    '''
    for sim_method in ["Euler", "EulerCromer", "EulerRichardson"]:
        for func in ["sin", "half_sin", "triangle", "rectangle"]:
            draw(sim_method, func)
